#include <os/mm.h>
#include <os/lock.h>
#include <os/time.h>
#include <os/irq.h>
#include <screen.h>
#include <stdio.h>
#include <assert.h>
#include <os/queue.h>
#include <os/sched.h>
#include <sbi.h>
//extern queue_t sleep_queue;
extern queue_t ready_queue;

extern void __global_pointer$();
pcb_t pcb[NUM_MAX_TASK];
int pid_status[NUM_MAX_TASK]={1};
const ptr_t cpu0_stack = INIT_KERNEL_STACK + PAGE_SIZE;
const ptr_t cpu1_stack = INIT_KERNEL_STACK + PAGE_SIZE + PAGE_SIZE;
pcb_t cpu0_pcb = {
    .pid = 0,
    .status = TASK_RUNNING,
    .kernel_sp = (ptr_t)cpu0_stack,
    .user_sp = (ptr_t)cpu0_stack,
    .preempt_count = 0,
    .kernel_context.regs[2]=cpu0_stack,
    .user_context.regs[2]=cpu0_stack,
    .mask = 0x1
};

pcb_t cpu1_pcb = {
    .pid = 0,
    .status = TASK_RUNNING,
    .kernel_sp = (ptr_t)cpu1_stack,
    .user_sp = (ptr_t)cpu1_stack,
    .preempt_count = 0,
    .kernel_context.regs[2]=cpu1_stack,
    .user_context.regs[2]=cpu1_stack,
    .mask = 0x2
};
/* current running task PCB */
pcb_t * volatile current_running[2];

/* global process id */
pid_t process_id = 1;
void timer_check()
{
    int current=get_timer();
    for(int i=1;i<NUM_MAX_TASK;i++){
        if(current >= pcb[i].ddl && pcb[i].status==TASK_SLEEPING){
            pcb[i].status = TASK_READY;
            pcb[i].ddl = 0;
            queue_push(&ready_queue,(void *)(&pcb[i]));

        }
    }
}



/*void scheduler(void)
{
    // Modify the current_running pointer.
    int time_elapsed = get_ticks();
    int rand=time_elapsed % 6;
    if(current_running[current_cpu]->status!=TASK_BLOCKED){
      current_running[current_cpu]->status = TASK_READY;
    if(current_running[current_cpu]->pid!=1)
        queue_push(current_queue,(void *)current_running[current_cpu]);
    }
    queue_t *old=current_queue;
    if( rand > 2)
      current_queue=&ready_queue[0];
    else if(rand > 0)
      current_queue=&ready_queue[1];
    else
      current_queue=&ready_queue[2];
    if(!queue_is_empty(current_queue))
      current_running[current_cpu]=(pcb_t *)queue_dequeue(current_queue);
    else
      current_queue = old;
    current_running[current_cpu]->status=TASK_RUNNING;
    vt100_move_cursor(current_running[current_cpu]->cursor_x,
                      current_running[current_cpu]->cursor_y);
}*/
extern unsigned long slave_cpu;
extern void wait_for_process();
extern void ret_from_exception();
void scheduler(void)
{
    // Modify the current_running pointer.
    //timer_check();
    int current_cpu=get_current_cpu_id();
    if(current_running[current_cpu]->pid == 0 ){
        current_running[current_cpu]->kernel_context.regs[1]=(uintptr_t)&ret_from_exception;
        current_running[current_cpu]->kernel_context.sepc = (uintptr_t)&wait_for_process;
        current_running[current_cpu]->user_context.regs[1]=(uintptr_t)&ret_from_exception;
        current_running[current_cpu]->user_context.sepc = (uintptr_t)&wait_for_process;
    }
    if(current_running[current_cpu]->status==TASK_KILLED)
        do_exit();
    if(current_running[current_cpu]->status != TASK_BLOCKED 
            && current_running[current_cpu]->status!=TASK_EXITED 
            && current_running[current_cpu]->status!=TASK_SLEEPING
            && current_running[current_cpu]->status!=TASK_KILLED){
            current_running[current_cpu]->status = TASK_READY;
            queue_push(&ready_queue,(void *)current_running[current_cpu]);
            current_running[current_cpu]->in_queue = &ready_queue;
        }
    pcb_t *try = current_running[current_cpu];
    while(!queue_is_empty(&ready_queue)){
		try=(pcb_t *)queue_dequeue(&ready_queue);
        if(try->mask== 3 || (try->mask == 1 && current_cpu== 0) || (try->mask == 2 && current_cpu == 1)){
            break;
        }else{
            queue_push(&ready_queue,(void *)try);      
        }
        try = current_running[current_cpu];
    }
    current_running[current_cpu] = try;
    current_running[current_cpu]->status=TASK_RUNNING;
    vt100_move_cursor(current_running[current_cpu]->cursor_x,
                      current_running[current_cpu]->cursor_y);
}
void do_sleep(uint32_t sleep_time)
{
    int current_cpu=get_current_cpu_id();
    if(sleep_time==0)
        return;
    current_running[current_cpu]->status = TASK_SLEEPING;
    current_running[current_cpu]->ddl = get_timer()+sleep_time;
    //queue_push(&sleep_queue,current_running[current_cpu]);
    do_scheduler();
    
}
int do_givecursor(int id){
    int current_cpu=get_current_cpu_id();
    if(id==0)
        return current_running[current_cpu]->cursor_x;
    else
        return current_running[current_cpu]->cursor_y;
}
void do_block(queue_t *queue)
{
    int current_cpu=get_current_cpu_id();
    // TODO: block the pcb task into the block queue
    current_running[current_cpu]->status=TASK_BLOCKED;
    queue_push(queue,(void *)current_running[current_cpu]);
    current_running[current_cpu]->in_queue = queue;
    do_scheduler();
}

void do_unblock(queue_t *queue)
{
    int current_cpu=get_current_cpu_id();
    // TODO: unblock the `pcb` from the block queue
    pcb_t *p;
    try_again:
    if(!queue_is_empty(queue)){
        p=(pcb_t *)queue_dequeue(queue);
        if(p->status==TASK_EXITED)
            goto try_again;
        p->status=TASK_READY;
        queue_push(&ready_queue,(void *)p);
        p->in_queue = &ready_queue;
    }
}

pid_t get_newpcb(){
    for (int i = 0; i < NUM_MAX_TASK; ++i)
    {
        if(pcb[i].status==TASK_EXITED)
            return i;
    }
}

pid_t do_spawn(task_info_t *task, void* arg, spawn_mode_t mode){
    int current_cpu=get_current_cpu_id();
    int i=get_newpcb();
    pcb[i].pid= process_id++;
    pcb[i].preempt_count=1;
    pcb[i].kernel_sp = pcb[i-1].kernel_sp+4096*2;
    pcb[i].user_sp = pcb[i-1].user_sp+4096*2;
    for(int j=0;j<32;j++)
        pcb[i].kernel_context.regs[j]=0;
    pcb[i].kernel_context.regs[1]=(uintptr_t)ret_from_exception;
    pcb[i].kernel_context.regs[2]=pcb[i].kernel_sp;
    pcb[i].kernel_context.regs[3]=__global_pointer$;
    pcb[i].kernel_context.regs[10]=arg;
    pcb[i].kernel_context.sstatus=34;
    queue_init(&pcb[i].wait_queue);
    pcb[i].kernel_context.sepc=(uintptr_t)task->entry_point;
    for(int j=0;j<32;j++)
        pcb[i].user_context.regs[j]=0;
    pcb[i].user_context.regs[1]=(uintptr_t)ret_from_exception;
    pcb[i].user_context.regs[2]=pcb[i].user_sp;
    pcb[i].user_context.regs[3]=__global_pointer$;
    pcb[i].user_context.regs[10]=arg;
    pcb[i].user_context.sstatus=34;
    pcb[i].user_context.sepc=(uintptr_t)task->entry_point;
    pcb[i].status = TASK_READY;
    pcb[i].type = task->type;
    queue_push(&ready_queue,(void *)&pcb[i]);
    pcb[i].in_queue = &ready_queue;
    pcb[i].cursor_x=0;
    pcb[i].cursor_y=0;
    pcb[i].mask = current_running[current_cpu]->mask;
    return pcb[i].pid;
}
void do_exit(void){
    int current_cpu=get_current_cpu_id();
    current_running[current_cpu]->status=TASK_EXITED;
    while(!queue_is_empty(&(current_running[current_cpu]->wait_queue))){
        do_unblock(&(current_running[current_cpu]->wait_queue));
    }
    while(!queue_is_empty(&(current_running[current_cpu]->lock_queue))){
        mutex_lock_t *lock=(mutex_lock_t*)queue_dequeue(&(current_running[current_cpu]->lock_queue));
        do_mutex_lock_release(lock);
    }
    do_scheduler();
}
void do_clear()
{
    screen_clear();
}
int do_kill(pid_t pid){
    int current_cpu=get_current_cpu_id();
    int i=0;
    for (i = 0; i < NUM_MAX_TASK; ++i)
    {
        if(pcb[i].pid==pid){
            break;
        }
    }
    if((current_running[0]->pid == pid && current_cpu == 1)|| current_running[1]->pid == pid&&current_cpu==0){
        pcb[i].status == TASK_KILLED;
        return 0;
    }
    pcb[i].status=TASK_EXITED;
    while(!queue_is_empty(&pcb[i].wait_queue)){
        do_unblock(&pcb[i].wait_queue);
    }
    while(!queue_is_empty(&(current_running[current_cpu]->lock_queue))){
        mutex_lock_t *lock=(mutex_lock_t*)queue_dequeue(&(current_running[current_cpu]->lock_queue));
        do_mutex_lock_release(lock);
    }
    if(&pcb[i]==current_running[current_cpu])
        do_scheduler();
}
int do_waitpid(pid_t pid){
    int current_cpu=get_current_cpu_id();
    int i=0;
    for (i = 0; i < NUM_MAX_TASK; ++i)
    {
        if(pcb[i].pid==pid)
            break;
    }
    if(pcb[i].status==TASK_EXITED)
        return 0;
    do_block(&(pcb[i].wait_queue));
    return 1;
}
void do_process_show(){
    int current_cpu=get_current_cpu_id();
    int i=0,j=0;
    for(j=0;j<process_id;j++){
        if(pcb[j].status==TASK_RUNNING){
            prints("[%d] PID: %d STATUS: TASK_RUNNING\n\r",i,pcb[j].pid);
            i++;
        }
        else if(pcb[j].status==TASK_BLOCKED&&pcb[j].pid!=0){
            prints("[%d] PID: %d STATUS: TASK_BLOCKED\n\r",i,pcb[j].pid);
            i++; 
        }
        else if(pcb[j].status==TASK_READY){
            prints("[%d] PID: %d STATUS: TASK_READY\n\r",i,pcb[j].pid);
            i++; 
        }
    }
}
pid_t do_getpid(){
    int current_cpu=get_current_cpu_id();
    return current_running[current_cpu]->pid;
}

void do_taskset(int mask,int pid){
    int i=0;
    for (i = 0; i < NUM_MAX_TASK; ++i)
    {
        if(pcb[i].pid==pid)
            break;
    }
    pcb[i].mask = mask;
}
